from ultralytics import YOLO
import numpy as np
import json

# 加载模型
model = YOLO("best.pt")

def recognize_image(image_path):
    # 预测图片
    results = model(image_path, show=True)
    names_dict = results[0].names
    probs = results[0].probs.data.tolist()

    # 返回识别结果
    result = {
        "prediction": names_dict[np.argmax(probs)]
    }
    return json.dumps(result, ensure_ascii=False)

if __name__ == "__main__":
    import sys
    image_path = sys.argv[1]
    print(recognize_image(image_path))